package seccomp

import (
	"errors"
	"fmt"
	"unsafe"

	"golang.org/x/sys/unix"
)

// BPFProgram 封装 BPF 程序
type BPFProgram struct {
	Filter []unix.SockFilter
}

// NewBPFProgram 创建新的 BPF 程序
func NewBPFProgram() *BPFProgram {
	return &BPFProgram{}
}

// Load 加载 Seccomp 过滤器
func (b *BPFProgram) Load() error {
	if err := b.Validate(); err != nil {
		return fmt.Errorf("seccomp validation failed: %w", err)
	}
	// 必须先设置 PR_SET_NO_NEW_PRIVS
	if err := unix.Prctl(unix.PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0); err != nil {
		return fmt.Errorf("prctl PR_SET_NO_NEW_PRIVS failed: %w", err)
	}

	// 准备 BPF 程序结构
	prog := &unix.SockFprog{
		Len:    uint16(len(b.Filter)),
		Filter: &b.Filter[0],
	}

	// 加载 Seccomp 过滤器
	_, _, errno := unix.Syscall(
		unix.SYS_SECCOMP,
		unix.SECCOMP_SET_MODE_FILTER,
		unix.SECCOMP_FILTER_FLAG_TSYNC, // 同步到所有线程
		uintptr(unsafe.Pointer(prog)),
	)
	if errno != 0 {
		return fmt.Errorf("seccomp load failed: %v", errno)
	}
	return nil
}

// AllowSyscall 允许特定系统调用
func (b *BPFProgram) AllowSyscall(syscallNum uint32) {
	/*
		BPF 指令序列:
		0: LD | W | ABS (0)  // 加载系统调用号
		1: JEQ syscallNum -> 允许 (jt=1, jf=0)
		2: RET KILL         // 默认拒绝
		3: RET ALLOW        // 允许执行
	*/
	b.Filter = append(b.Filter,
		// 加载系统调用号到累加器
		unix.SockFilter{
			Code: unix.BPF_LD | unix.BPF_W | unix.BPF_ABS,
			K:    0, // seccomp_data.nr 的偏移量
		},
		// 比较系统调用号
		unix.SockFilter{
			Code: unix.BPF_JMP | unix.BPF_JEQ | unix.BPF_K,
			Jt:   1, // 匹配时跳转到允许指令
			Jf:   0, // 不匹配时继续
			K:    syscallNum,
		},
		// 默认拒绝 (KILL)
		unix.SockFilter{
			Code: unix.BPF_RET | unix.BPF_K,
			K:    unix.SECCOMP_RET_KILL_PROCESS,
		},
		// 允许执行 (ALLOW)
		unix.SockFilter{
			Code: unix.BPF_RET | unix.BPF_K,
			K:    unix.SECCOMP_RET_ALLOW,
		},
	)
}

// BlockSyscall 阻止特定系统调用
func (b *BPFProgram) BlockSyscall(syscallNum uint32) {
	/*
		BPF 指令序列:
		0: LD | W | ABS (0)  // 加载系统调用号
		1: JEQ syscallNum -> 拒绝 (jt=0, jf=1)
		2: RET ALLOW         // 默认允许
		3: RET KILL          // 拒绝执行
	*/
	b.Filter = append(b.Filter,
		// 加载系统调用号到累加器
		unix.SockFilter{
			Code: unix.BPF_LD | unix.BPF_W | unix.BPF_ABS,
			K:    0, // seccomp_data.nr 的偏移量
		},
		// 比较系统调用号
		unix.SockFilter{
			Code: unix.BPF_JMP | unix.BPF_JEQ | unix.BPF_K,
			Jt:   0, // 匹配时跳转到拒绝指令
			Jf:   1, // 不匹配时继续
			K:    syscallNum,
		},
		// 拒绝执行 (KILL)
		unix.SockFilter{
			Code: unix.BPF_RET | unix.BPF_K,
			K:    unix.SECCOMP_RET_KILL_PROCESS,
		},
		// 默认允许 (ALLOW)
		unix.SockFilter{
			Code: unix.BPF_RET | unix.BPF_K,
			K:    unix.SECCOMP_RET_ALLOW,
		},
	)
}

// DefaultSandbox 创建默认沙箱规则
func DefaultSandbox() *BPFProgram {
	prog := NewBPFProgram()

	// 允许基本系统调用
	allowList := []uint32{
		unix.SYS_READ,
		unix.SYS_WRITE,
		unix.SYS_CLOSE,
		unix.SYS_FSTAT,
		unix.SYS_EXIT,
		unix.SYS_EXIT_GROUP,
		unix.SYS_MMAP,
		unix.SYS_MUNMAP,
		unix.SYS_MPROTECT,
		unix.SYS_BRK,
		unix.SYS_FACCESSAT,
		unix.SYS_RT_SIGRETURN,
	}

	for _, sc := range allowList {
		prog.AllowSyscall(sc)
	}

	return prog
}

// OptimizedSandbox 优化版沙箱规则（使用跳转表）
func OptimizedSandbox() *BPFProgram {
	prog := NewBPFProgram()

	// 基本系统调用白名单
	allowList := []uint32{
		unix.SYS_READ,
		unix.SYS_WRITE,
		unix.SYS_CLOSE,
		unix.SYS_FSTAT,
		unix.SYS_EXIT,
		unix.SYS_EXIT_GROUP,
		unix.SYS_MMAP,
		unix.SYS_MUNMAP,
		unix.SYS_MPROTECT,
		unix.SYS_BRK,
		unix.SYS_FACCESSAT,
		unix.SYS_RT_SIGRETURN,
	}

	// 初始指令：加载系统调用号
	prog.Filter = append(prog.Filter, unix.SockFilter{
		Code: unix.BPF_LD | unix.BPF_W | unix.BPF_ABS,
		K:    0, // seccomp_data.nr
	})

	// 为每个系统调用生成检查指令
	for i, sc := range allowList {
		prog.Filter = append(prog.Filter, unix.SockFilter{
			Code: unix.BPF_JMP | unix.BPF_JEQ | unix.BPF_K,
			Jt:   1,  // 匹配时跳转1步
			Jf:   0,  // 不匹配继续
			K:    sc, // 系统调用号
		})

		// 跳转到允许执行
		prog.Filter = append(prog.Filter, unix.SockFilter{
			Code: unix.BPF_JMP | unix.BPF_JA,
			Jt:   uint8(len(allowList) - i + 1), // 跳转到ALLOW
			Jf:   0,
			K:    0,
		})
	}

	// 默认拒绝
	prog.Filter = append(prog.Filter, unix.SockFilter{
		Code: unix.BPF_RET | unix.BPF_K,
		K:    unix.SECCOMP_RET_KILL_PROCESS,
	})

	// 允许执行
	prog.Filter = append(prog.Filter, unix.SockFilter{
		Code: unix.BPF_RET | unix.BPF_K,
		K:    unix.SECCOMP_RET_ALLOW,
	})

	return prog
}

// ApplyToPid 为指定PID应用Seccomp过滤器
func (b *BPFProgram) ApplyToPid(pid int) error {
	// 使用ptrace附加到目标进程
	if err := unix.PtraceAttach(pid); err != nil {
		return fmt.Errorf("ptrace attach failed: %w", err)
	}
	defer unix.PtraceDetach(pid)

	// 等待进程停止
	var status unix.WaitStatus
	_, err := unix.Wait4(pid, &status, 0, nil)
	if err != nil {
		return fmt.Errorf("wait4 failed: %w", err)
	}

	// 加载Seccomp过滤器
	if err := b.Load(); err != nil {
		return fmt.Errorf("seccomp load failed: %w", err)
	}

	return nil
}

// ApplyToCurrentProcess 为当前进程应用Seccomp过滤器
func (b *BPFProgram) ApplyToCurrentProcess() error {
	return b.Load()
}

func (b *BPFProgram) Validate() error {
	if len(b.Filter) == 0 {
		return errors.New("empty seccomp filter")
	}

	// 确保最后一条指令是返回指令
	lastInstruction := b.Filter[len(b.Filter)-1]
	if lastInstruction.Code&unix.BPF_RET == 0 {
		return errors.New("last instruction must be a return")
	}

	// 确保没有危险指令
	for _, inst := range b.Filter {
		if inst.Code&unix.BPF_ST == unix.BPF_ST {
			return errors.New("BPF_ST instruction not allowed")
		}
	}

	return nil
}
